import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding

class AESCipher:
    """
    AES加密解密类
    """
    def __init__(self, key, iv):
        """
        初始化AES加密解密类
        :param key: AES密钥，长度为16、24或32字节
        :param iv: 初始化向量，长度为16字节
        """
        self.key = key
        self.iv = iv
        self.backend = default_backend()
        self.cipher = Cipher(algorithms.AES(self.key), modes.CBC(self.iv), backend=self.backend)
        self.encryptor = self.cipher.encryptor()
        self.decryptor = self.cipher.decryptor()

    def encrypt(self, data):
        """
        AES加密
        :param data: 待加密的数据
        :return: 加密后的数据
        """
        padder = padding.PKCS7(algorithms.AES.block_size).padder()
        data = padder.update(data) + padder.finalize()
        return self.encryptor.update(data) + self.encryptor.finalize()
    
    def encrypt_from_str(self, data):
        """
        AES加密
        :param data: 待加密的数据
        :return: 加密后的数据
        """
        # 将字符串转换为字节数组
        # data = bytes.fromhex(data)
        data = data.encode("utf-8")
        padder = padding.PKCS7(algorithms.AES.block_size).padder()
        data = padder.update(data) + padder.finalize()
        return self.encryptor.update(data) + self.encryptor.finalize()
    
    def decrypt(self, data):
        """
        AES解密
        :param data: 待解密的数据
        :return: 解密后的数据
        """
        unpadded_data = self.decryptor.update(data) + self.decryptor.finalize()
        return AESCipher.__remove_padding(unpadded_data)
    
    def decrypt_to_str(self, data):
        """
        AES解密
        :param data: 待解密的数据
        :return: 解密后的数据
        """
        unpadded_data = self.decryptor.update(data) + self.decryptor.finalize()
        return AESCipher.__remove_padding(unpadded_data)
    
    def generate_aes_key_iv():
        """生成AES-256的key和IV"""
        key = os.urandom(32).upper() # AES-256需要32字节的Key
        iv = os.urandom(16).upper()  # IV通常为16字节
        return key, iv

    def createKeyAndIv(count):
        """根据给定的数量生成多组AES的key和iv，并打印出来"""
        for i in range(count):
            key, iv = AESCipher.generate_aes_key_iv()
            print(f"Key {i+1}: {key.hex()}")
            print(f"IV {i+1}: {iv.hex()}")
            print("---")

    def __remove_padding(padded_data):
        """移除给定数据的PKCS7填充"""
        unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
        data = unpadder.update(padded_data)
        return data + unpadder.finalize()
